from flask import Flask, request
import subprocess

app = Flask(__name__)

sh_file = "sh_file_path"

py_file = "baichuan2/run_baichuan2.py"

config_file = "baichuan2/run_baichuan2_13b.yaml"

load_checkpoint = "model_dir"

rank_table_file = "rank_table_file_path"


@app.route('/get_result', methods=['GET'])
def get_keyword():
    predict_data = request.args.get('keyword')

    arg1 = "\"python " + py_file + " --config " + config_file + " --run_mode predict" \
           + " --use_parallel True" + " --load_checkpoint " + load_checkpoint \
           + " --auto_trans_ckpt True" + " --predict_data " + predict_data + "\""
    arg3 = [0, 2] #不行的话尝试 "[0,2]" ，用双引号括起来
    args = ["sh", sh_file, arg1, rank_table_file, arg3, 2]
    try:
        result = subprocess.run(args,
                                text=True,
                                capture_output=True,
                                encoding="utf-8")
        if result.returncode == 0:
            return result.stdout
        else:
            return result.stderr

    except Exception as e:
        return "脚本执行错误" + e


if __name__ == '__main__':
    app.run(debug=True, port=9002)
